# from .vqvae import (
#     VQVAEConfiguration,
#     VQVAEModel,
#     VQVAETrainer,
#     VQVAEModelWrapper
# )
# from .causal_vqvae import (
#     CausalVQVAEConfiguration,
#     CausalVQVAETrainer,
#     CausalVQVAEModel, CausalVQVAEModelWrapper
# )
from .causal_vae import CausalVAEModelWrapper
from .causal_vae.modeling_causalvae import CausalVAEModel
from .ema_model import EMA

videobase_ae_stride = {
    "CausalVAEModel_4x8x8": [4, 8, 8],
    # 'CausalVQVAEModel_4x4x4': [4, 4, 4],
    # 'CausalVQVAEModel_4x8x8': [4, 8, 8],
    # 'VQVAEModel_4x4x4': [4, 4, 4],
    # 'OpenVQVAEModel_4x4x4': [4, 4, 4],
    # 'VQVAEModel_4x8x8': [4, 8, 8],
    # 'bair_stride4x2x2': [4, 2, 2],
    # 'ucf101_stride4x4x4': [4, 4, 4],
    # 'kinetics_stride4x4x4': [4, 4, 4],
    # 'kinetics_stride2x4x4': [2, 4, 4],
}

videobase_ae_channel = {
    "CausalVAEModel_4x8x8": 4,
    # 'CausalVQVAEModel_4x4x4': 4,
    # 'CausalVQVAEModel_4x8x8': 4,
    # 'VQVAEModel_4x4x4': 4,
    # 'OpenVQVAEModel_4x4x4': 4,
    # 'VQVAEModel_4x8x8': 4,
    # 'bair_stride4x2x2': 256,
    # 'ucf101_stride4x4x4': 256,
    # 'kinetics_stride4x4x4': 256,
    # 'kinetics_stride2x4x4': 256,
}

videobase_ae = {
    "CausalVAEModel_4x8x8": CausalVAEModelWrapper,
    # 'CausalVQVAEModel_4x4x4': CausalVQVAEModelWrapper,
    # 'CausalVQVAEModel_4x8x8': CausalVQVAEModelWrapper,
    # 'VQVAEModel_4x4x4': VQVAEModelWrapper,
    # 'VQVAEModel_4x8x8': VQVAEModelWrapper,
    # "bair_stride4x2x2": VQVAEModelWrapper,
    # "ucf101_stride4x4x4": VQVAEModelWrapper,
    # "kinetics_stride4x4x4": VQVAEModelWrapper,
    # "kinetics_stride2x4x4": VQVAEModelWrapper,
}
